#!/usr/bin/R

#set params/initial conditions
n.ponds=3 #number of ponds
area=c(100,100,100) #surface areas (m^2)
bottom.elev=c(0,0,0) #bottom elevations above datum (m)
weir.elev=c(1,1,1) #elevation of weirs above datum (m)
weir.width=c(1,1,1) #width of weirs (m)
sea.elev=1 #elevation on sea-side of final weir
t.final=100 #simulation length (hrs)
dt=0.1 #timestep length(hrs)
n.timesteps=as.integer(t.final/dt) #number of timesteps
q=array(0,dim=c(n.ponds,n.timesteps)) #initialize flows over weirs (cms)
q.temp=array(0,dim=n.ponds) #temporary variable for Runge Kutta routine
inflow=array(1,dim=n.timesteps) #inflow to first pond for all times (cms)
#note: h[1:n.ponds,t]=pond elevations; h[n.ponds+1,t]=sea surface elevation 
h=array(0,dim=c(n.ponds+1,n.timesteps)) #initialize pond levels (m)
h[1:n.ponds,1]=c(1,1,1) #initial pond elevations above datum (m)
h[n.ponds+1,]=1 #set sea-surface elev above datum for all times (m)
K1=K2=K3=array(0,dim=n.ponds) #temporary variables for Runge Kutta routine

#generate color array
library(fields)
my.colors=array(tim.colors(),dim=(n.timesteps+1))

#plot initial conditions
plot(type='l',h[,1],col=my.colors[1])

#loop through time
for (t in 1:(n.timesteps-1)) {
	#determine which weirs are submerged
	submerged=forward.flow=backward.flow=noflow=array(FALSE,dim=n.ponds) #true if outflow weir is submerged
	submerged[which(h[1:n.ponds,t]>=weir.elev & h[2:(n.ponds+1),t]>=weir.elev)] = TRUE
	forward.flow[which(h[1:n.ponds,t]>=weir.elev & h[2:(n.ponds+1),t]<weir.elev)] = TRUE
	backward.flow[which(h[1:n.ponds,t]<weir.elev & h[2:(n.ponds+1),t]>=weir.elev)] = TRUE
	noflow[which(h[1:n.ponds,t]<weir.elev & h[2:(n.ponds+1),t]<weir.elev)] = TRUE
	i=which(submerged==TRUE)
	j=which(forward.flow==TRUE)
	k=which(backward.flow==TRUE)
	#determine outflow
		q[i,t] = 3.33*weir.width[i]*sqrt(abs(h[i,t]-h[i+1,t]))*
			((h[i,t]-weir.elev[i])+(h[i+1,t]-weir.elev[i])*.381)
		nonzero=which(submerged==TRUE & (q[,t]!=0))
		q[nonzero,t] = q[nonzero,t]*(h[nonzero,t]-h[nonzero+1,t])/
			abs(h[nonzero,t]-h[nonzero+1,t])
		q[j,t]=3.33*weir.width[j]*((h[j,t]-weir.elev[j])^(3/2))
		q[k,t]=(-1)*3.33*weir.width[k]*((h[k+1,t]-weir.elev[k])^(3/2))
	#determine new pond levels
		#first pond
		K1[1]=dt/area[1]*(inflow[t]-q[1,t])
		#subsequent ponds
		K1[2:n.ponds]=dt/area[2:n.ponds]*(q[1:(n.ponds-1),t]-q[2:n.ponds,t])
			#intermediates
			h.temp=c(h[1:n.ponds,t]+(1/3)*K1,h[n.ponds+1,t])
			q.temp[i]=3.33*weir.width[i]*sqrt(abs(h.temp[i]-h.temp[i+1]))*
				((h.temp[i]-weir.elev[i])+(h.temp[i+1]-weir.elev[i])*.381)
			nonzero=which(submerged==TRUE & (q.temp!=0))
			q.temp[nonzero] = q.temp[nonzero]*(h.temp[nonzero]-h.temp[nonzero+1])/
				abs(h.temp[nonzero]-h.temp[nonzero+1])
			q.temp[j]=3.33*weir.width[j]*((h.temp[j]-weir.elev[j])^(3/2))
			q.temp[k]=(-1)*3.33*weir.width[k]*((h.temp[k+1]-weir.elev[k])^(3/2))
		#first pond
		K2[1]=dt/area[1]*(inflow[t]-q.temp[1])
		#subsequent ponds
		K2[2:n.ponds]=dt/area[2:n.ponds]*(q.temp[1:(n.ponds-1)]-q.temp[2:n.ponds])
			#intermediates
                        h.temp=c(h[1:n.ponds,t]+(2/3)*K2,h[n.ponds+1,t])
			q.temp[i]=3.33*weir.width[i]*sqrt(abs(h.temp[i]-h.temp[i+1]))*
                                ((h.temp[i]-weir.elev[i])+(h.temp[i+1]-weir.elev[i])*.381)
                        nonzero=which(submerged==TRUE & (q.temp!=0))
                        q.temp[nonzero] = q.temp[nonzero]*(h.temp[nonzero]-h.temp[nonzero+1])/
                                abs(h.temp[nonzero]-h.temp[nonzero+1])
                        q.temp[j]=3.33*weir.width[j]*((h.temp[j]-weir.elev[j])^(3/2))
                        q.temp[k]=(-1)*3.33*weir.width[k]*((h.temp[k+1]-weir.elev[k])^(3/2))
		#first pond
		K3[1]=dt/area[1]*(inflow[t]-q.temp[1])
		#subsequent ponds
		K3[2:n.ponds]=dt/area[2:n.ponds]*(q.temp[1:(n.ponds-1)]-q.temp[2:n.ponds])
		h[1:n.ponds,t+1]=h[1:n.ponds,t]+(1/4)*K1+(3/4)*K3
	#determine outflow @ end of timestep
		#determine which weirs are submerged
        	submerged=forward.flow=backward.flow=noflow=array(FALSE,dim=n.ponds) #true if outflow weir is submerged
        	submerged[which(h[1:n.ponds,t+1]>=weir.elev & h[2:(n.ponds+1),t+1]>=weir.elev)] = TRUE
        	forward.flow[which(h[1:n.ponds,t+1]>=weir.elev & h[2:(n.ponds+1),t+1]<weir.elev)] = TRUE
        	backward.flow[which(h[1:n.ponds,t+1]<weir.elev & h[2:(n.ponds+1),t+1]>=weir.elev)] = TRUE
        	noflow[which(h[1:n.ponds,t+1]<weir.elev & h[2:(n.ponds+1),t+1]<weir.elev)] = TRUE
        	i=which(submerged==TRUE)
        	j=which(forward.flow==TRUE)
        	k=which(backward.flow==TRUE)
		#determine new outflows
		q.temp[i]=3.33*weir.width[i]*sqrt(abs(h[i,t+1]-h[i+1,t+1]))*
                	((h[i,t+1]-weir.elev[i])+(h[i+1,t+1]-weir.elev[i])*.381)
                nonzero=which(submerged==TRUE & (q.temp!=0))
                q.temp[nonzero] = q.temp[nonzero]*(h[nonzero,t+1]-h[nonzero+1,t+1])/
                        abs(h[nonzero,t+1]-h[nonzero+1,t+1])
                q.temp[j]=3.33*weir.width[j]*((h[j,t+1]-weir.elev[j])^(3/2))
                q.temp[k]=(-1)*3.33*weir.width[k]*((h[k+1,t+1]-weir.elev[k])^(3/2))
	#if 'new' outflow is very different....???
		#DO SOMETHING HERE#
	lines(h[,t+1],col=my.colors[t+1])
}
